#include "ag.h"
#include <rtl/hash.h>
#include <rtl/print.h>
#include <rtl/memory.h>

#include "bootinfo.h"

#define MOD "PATCH"

VOID AgpLkPanicHandler()
{
    // disable watchdog to prevent reboot
    *((VOLATILE UINT32 *)0x10007000) = 0x22000000;
    for (;;)
        ;
}

VOID AgpPatchLgK102017M20ENougatLk(PVOID Base, SIZE Size, UINT32 BootMode, UINT32 Flags)
{
    if (BootMode == FASTBOOT)
    {
        ((VOLATILE UINT16 *)Base)[0x24BB0 / 2] = 0xe023;
        RtlPrintf("[%s] Forced fastboot mode\n", MOD);
    }

    /*
    // Jump on panic
    ((VOLATILE UINT16 *)Base)[0x2DC6 / 2] = 0x4770;
    ((VOLATILE UINT16 *)Base)[0x2D3BA / 2] = 0xbf00;
    ((VOLATILE UINT32 *)Base)[0x2D3BC / 4] = 0xf000f8df;
    ((VOLATILE UINT16 *)Base)[0x2D3B4 / 2] = 0xbf00;
    ((VOLATILE UINT16 *)Base)[0x2D3B2 / 2] = 0xbf00;

    // JUMP VECTOR
    ((VOLATILE UINT32 *)Base)[0x2D3C0 / 4] = (UINT32)AgpLkPanicHandler;

    RtlPrintf("[%s] LK will return to agent on panic (jump to %p).\n", MOD, AgpLkPanicHandler);*/
}

typedef struct _MAP
{
    UINT8 Sha256Sum[SHA256_DIGEST_SIZE];
    VOID(*P)
    (PVOID Base, SIZE Size, UINT32 BootMode, UINT32 Flags);
} MAP, *PMAP;

STATIC MAP Patches[] = {
    {{0x95, 0xa7, 0xc6, 0xb1, 0xe7, 0x44, 0x5c, 0xb8, 0x65, 0xb0, 0xde, 0xd9, 0xe3, 0xe3, 0x04, 0xbf, 0xd5, 0x94, 0x58, 0x39, 0xb1, 0xdb, 0xf7, 0xdc, 0x8a, 0x6f, 0xaf, 0x42, 0xea, 0xbf, 0xd3, 0xc2},
     AgpPatchLgK102017M20ENougatLk}, // LK K10 2017 M250E with nougat

    /*// Null terminate
    {{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
     NULL}*/
};

VOID AgpPatchLk(PVOID Base, PVOID CodeStart, SIZE Size, UINT32 BootMode, UINT32 Flags)
{
    CHAR Digest[SHA256_DIGEST_SIZE];
    RTL_SHA256_CONTEXT Sha;
    RtlSha256Initialize(&Sha);
    RtlSha256Update(&Sha, Base, Size);
    RtlSha256Finish(&Sha, Digest);

    RtlPrintf(
        "[%s] LK size=0x%08x SHA256=",
        MOD,
        Size);
    for (int i = 0; i < SHA256_DIGEST_SIZE; i++)
        RtlPrintf("%02x", Digest[i]);
    RtlPrintf("\n");

    BOOLEAN PatchFound = FALSE;
    for (INT I = 0; I < ARRAY_SIZE(Patches); I++)
    {
        if (RtlCompareMemory(Digest, Patches[I].Sha256Sum, SHA256_DIGEST_SIZE) == 0)
        {
            if (!Patches[I].P)
                continue;
            Patches[I].P(CodeStart, Size, BootMode, Flags);
            PatchFound = TRUE;
            break;
        }
    }

    if (!PatchFound)
        RtlPrintf("[%s] No patch found.\n", MOD);
}